import os
import time
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
matplotlib.rcParams['font.family'] = ['Microsoft YaHei', 'SimHei', 'sans-serif']
matplotlib.rcParams['axes.unicode_minus'] = False
from tqdm import tqdm
from colorama import Fore, Back, Style, init
from tabulate import tabulate
from datetime import datetime
import psutil
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import shutil
import warnings
import gc
import json

# 忽略特定警告
warnings.filterwarnings('ignore', category=UserWarning)
warnings.filterwarnings('ignore', category=FutureWarning)

# 设置环境变量以禁用albumentations的版本检查
os.environ['ALBUMENTATIONS_VERSION_COMPATIBILITY_CHECK'] = '1'

# 初始化 colorama
init(autoreset=True)

class PrettyTrainer:
    """美观的模型训练输出管理器"""
    
    def __init__(self, model_name, save_dir="./runs"):
        """初始化训练器"""
        self.model_name = model_name
        self.epochs = 30  # 默认训练轮数
        
        # 清理旧的runs目录
        if os.path.exists(save_dir):
            try:
                shutil.rmtree(save_dir)
                print(f"{Fore.YELLOW}已清理旧的训练记录目录{Style.RESET_ALL}")
            except Exception as e:
                print(f"{Fore.RED}清理目录失败: {str(e)}{Style.RESET_ALL}")
        
        # 创建保存目录
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.save_dir = os.path.join(save_dir, f"{model_name}_{timestamp}")
        os.makedirs(self.save_dir, exist_ok=True)
        
        # 创建model目录用于保存最佳模型
        self.model_dir = "./model"
        os.makedirs(self.model_dir, exist_ok=True)
        
        # 创建或更新latest目录
        self.latest_dir = os.path.join(save_dir, "latest")
        if os.path.exists(self.latest_dir):
            shutil.rmtree(self.latest_dir)
        shutil.copytree(self.save_dir, self.latest_dir)
        
        # 初始化colorama
        init(autoreset=True)
        
        # 配置matplotlib中文显示
        import matplotlib.font_manager as fm
        # 查找系统中的中文字体
        font_paths = [
            'C:/Windows/Fonts/simhei.ttf',  # 黑体
            'C:/Windows/Fonts/msyh.ttc',    # 微软雅黑
            'C:/Windows/Fonts/simsun.ttc',  # 宋体
        ]
        
        font_found = False
        for font_path in font_paths:
            if os.path.exists(font_path):
                plt.rcParams['font.family'] = ['sans-serif']
                plt.rcParams['font.sans-serif'] = [fm.FontProperties(fname=font_path).get_name()]
                plt.rcParams['axes.unicode_minus'] = False
                print(f"{Fore.GREEN}成功加载中文字体: {font_path}{Style.RESET_ALL}")
                font_found = True
                break
                
        if not font_found:
            print(f"{Fore.YELLOW}警告：未找到中文字体，图表中文可能显示为方块{Style.RESET_ALL}")
        
        # 打印训练开始信息
        self._print_header()
        self._print_system_info()
        
        self.start_time = time.time()
        self.plot_interval = 5
        self.metrics_history = {
            'train_loss': [],
            'val_loss': [],
            'train_acc': [],
            'val_acc': [],
            'learning_rate': [],
            'epoch_time': []
        }
        self.best_val_acc = 0
        self.best_epoch = 0
        self.early_stop_counter = 0
        
    def _print_header(self):
        """打印训练开始的标题"""
        print("\n" + "="*80)
        print(f"{Fore.CYAN}{Style.BRIGHT}✨ 深度学习训练任务 ✨{Style.RESET_ALL}".center(90))
        print(f"{Fore.YELLOW}模型: {Style.BRIGHT}{self.model_name}{Style.RESET_ALL}".center(80))
        print(f"{Fore.YELLOW}保存路径: {Style.BRIGHT}{self.save_dir}{Style.RESET_ALL}".center(90))
        print("="*80 + "\n")
    
    def _print_system_info(self):
        """打印系统资源信息"""
        print(f"\n{Fore.CYAN}🖥️ 系统资源信息 🖥️{Style.RESET_ALL}\n")
        
        # 获取系统信息
        cpu_percent = psutil.cpu_percent()
        memory = psutil.virtual_memory()
        
        # 获取GPU信息
        gpu_info = []
        if torch.cuda.is_available():
            gpu_name = torch.cuda.get_device_name(0)
            gpu_memory = torch.cuda.get_device_properties(0).total_memory / (1024**3)
            gpu_memory_allocated = torch.cuda.memory_allocated(0) / (1024**3)
            gpu_memory_percent = (gpu_memory_allocated / gpu_memory) * 100
            
            gpu_info = [
                ["GPU设备", gpu_name],
                ["GPU显存总量", f"{gpu_memory:.2f} GB"],
                ["GPU显存使用", f"{gpu_memory_allocated:.2f} GB"],
                ["GPU显存占用率", f"{gpu_memory_percent:.2f}%"]
            ]
        
        # 创建系统信息表格
        system_info = [
            ["CPU使用率", f"{cpu_percent}%"],
            ["内存使用率", f"{memory.percent}%"],
            ["可用内存", f"{memory.available/(1024**3):.2f} GB"],
            ["总内存", f"{memory.total/(1024**3):.2f} GB"]
        ]
        
        if gpu_info:
            system_info.extend(gpu_info)
        
        print(tabulate(system_info, headers=["指标", "数值"], tablefmt="fancy_grid"))
    
    def process_data(self, description, total_steps):
        """显示数据处理进度"""
        print(f"\n>>> 数据处理: {description}")
        return tqdm(total=total_steps, 
                   desc="处理进度",
                   ncols=100,
                   bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]')
    
    def update_data_progress(self, pbar, step, message=None):
        """更新数据处理进度"""
        if message:
            pbar.set_postfix_str(message)
        pbar.update(step)
    
    def data_summary(self, data_info):
        """显示数据集信息摘要"""
        print(f"\n{Fore.CYAN}📊 数据集信息 📊{Style.RESET_ALL}\n")
        
        # 基本信息
        basic_info = [
            ["总图片数量", data_info["总图片数量"]],
            ["类别数量", data_info["类别数量"]],
            ["训练集比例", data_info["训练集比例"]],
            ["验证集比例", data_info["验证集比例"]],
            ["测试集比例", data_info["测试集比例"]]
        ]
        
        print(tabulate(basic_info, headers=["指标", "数值"], tablefmt="fancy_grid"))
        
        # 类别分布信息
        print(f"\n{Fore.YELLOW}类别分布:{Style.RESET_ALL}")
        class_dist = [[k, v] for k, v in data_info["class_distribution"].items()]
        print(tabulate(class_dist, headers=["类别", "图片数量"], tablefmt="simple"))
    
    def model_summary(self, model_info):
        """显示模型信息摘要"""
        print(f"\n{Fore.CYAN}🤖 模型信息 🤖{Style.RESET_ALL}\n")
        
        info_table = [
            ["模型类型", model_info["model_type"]],
            ["总参数量", f"{model_info['total_params']:,}"],
            ["可训练参数量", f"{model_info['trainable_params']:,}"],
            ["优化器", model_info["optimizer"]],
            ["初始学习率", model_info["learning_rate"]],
            ["损失函数", model_info["loss_function"]]
        ]
        
        print(tabulate(info_table, headers=["指标", "数值"], tablefmt="fancy_grid"))
    
    def start_epoch(self, epoch, learning_rate=None):
        """开始新的训练轮次"""
        print(f"\n{Fore.CYAN}🔄 第 {epoch}/{self.epochs} 轮训练开始 🔄{Style.RESET_ALL}")
        if learning_rate is not None:
            print(f"当前学习率: {learning_rate:.6f}")
        
        # 显示GPU内存使用情况
        if torch.cuda.is_available():
            allocated = torch.cuda.memory_allocated() / 1024**3
            reserved = torch.cuda.memory_reserved() / 1024**3
            print(f"GPU内存: 已分配 {allocated:.2f} GB, 已保留 {reserved:.2f} GB")
    
    def update_gpu_stats(self, allocated, reserved):
        """更新GPU统计信息"""
        if hasattr(self, '_gpu_stats_line'):
            print(f"\rGPU内存: 已分配 {allocated:.2f} GB, 已保留 {reserved:.2f} GB", end='')
        else:
            self._gpu_stats_line = True
    
    def end_epoch(self, metrics):
        """结束训练轮次，显示训练结果"""
        print(f"\n{Fore.GREEN}📊 训练结果 📊{Style.RESET_ALL}")
        
        table = [
            ["训练损失", f"{metrics['train_loss']:.4f}"],
            ["训练准确率", f"{metrics['train_acc']:.2f}%"],
            ["验证损失", f"{metrics['val_loss']:.4f}"],
            ["验证准确率", f"{metrics['val_acc']:.2f}%"],
            ["最佳准确率", f"{metrics['best_acc']:.2f}%"]
        ]
        
        print(tabulate(table, headers=["指标", "数值"], tablefmt="fancy_grid"))
        
        # 保存训练历史
        if not hasattr(self, 'history'):
            self.history = {k: [] for k in metrics.keys()}
        
        for k, v in metrics.items():
            self.history[k].append(v)
        
        # 绘制训练曲线
        self._plot_training_curves()
    
    def _plot_training_curves(self):
        """绘制训练曲线"""
        plt.figure(figsize=(12, 5))
        
        # 损失曲线
        plt.subplot(1, 2, 1)
        plt.plot(self.history['train_loss'], label='训练损失')
        plt.plot(self.history['val_loss'], label='验证损失')
        plt.title('损失曲线')
        plt.xlabel('轮次')
        plt.ylabel('损失')
        plt.legend()
        
        # 准确率曲线
        plt.subplot(1, 2, 2)
        plt.plot(self.history['train_acc'], label='训练准确率')
        plt.plot(self.history['val_acc'], label='验证准确率')
        plt.title('准确率曲线')
        plt.xlabel('轮次')
        plt.ylabel('准确率 (%)')
        plt.legend()
        
        plt.tight_layout()
        plt.savefig(os.path.join(self.save_dir, 'training_curves.png'))
        plt.close()
    
    def update_metrics(self, epoch, metrics, save_plot=False):
        """更新并显示训练指标"""
        for key, value in metrics.items():
            if key in self.metrics_history:
                self.metrics_history[key].append(value)
        
        if metrics.get('val_acc', 0) > self.best_val_acc:
            self.best_val_acc = metrics.get('val_acc', 0)
            self.best_epoch = epoch
            improved = f"{Fore.GREEN}✓ 最佳模型!{Style.RESET_ALL}"
            self.early_stop_counter = 0
        else:
            improved = ""
            self.early_stop_counter += 1
        
        epoch_data = [
            ["训练损失", f"{metrics.get('train_loss', 0):.4f}"],
            ["验证损失", f"{metrics.get('val_loss', 0):.4f}"],
            ["训练准确率", f"{metrics.get('train_acc', 0)*100:.2f}%"],
            ["验证准确率", f"{metrics.get('val_acc', 0)*100:.2f}%"]
        ]
        
        print(f"\n{Fore.GREEN}{Style.BRIGHT}📈 第 {epoch+1} 轮指标 📈{Style.RESET_ALL} {improved}")
        print(tabulate(epoch_data, headers=["指标", "数值"], tablefmt="fancy_grid"))
        print(f"{Fore.YELLOW}最佳验证准确率: {self.best_val_acc*100:.2f}% (轮次 {self.best_epoch+1}){Style.RESET_ALL}")
        
        if save_plot and (epoch + 1) % self.plot_interval == 0:
            self.plot_metrics()
    
    def plot_metrics(self):
        """绘制训练指标图表"""
        plt.figure(figsize=(12, 8))
        plt.subplot(2, 1, 1)
        plt.plot(self.metrics_history['train_loss'], label='训练损失')
        plt.plot(self.metrics_history['val_loss'], label='验证损失')
        plt.title('损失曲线')
        plt.xlabel('轮次')
        plt.ylabel('损失')
        plt.legend()
        
        plt.subplot(2, 1, 2)
        plt.plot(self.metrics_history['train_acc'], label='训练准确率')
        plt.plot(self.metrics_history['val_acc'], label='验证准确率')
        plt.title('准确率曲线')
        plt.xlabel('轮次')
        plt.ylabel('准确率')
        plt.legend()
        
        plt.tight_layout()
        plt.savefig(os.path.join(self.save_dir, f"metrics_epoch_{len(self.metrics_history['train_loss'])}.png"))
        plt.close()
    
    def finish_training(self):
        """完成训练，显示总结信息"""
        total_time = time.time() - self.start_time
        hours = int(total_time // 3600)
        minutes = int((total_time % 3600) // 60)
        seconds = int(total_time % 60)
        
        print(f"\n{Fore.GREEN}{Style.BRIGHT}🎉 训练完成! 🎉{Style.RESET_ALL}")
        print(f"总用时: {hours}小时 {minutes}分钟 {seconds}秒")
        print(f"最佳验证准确率: {self.best_val_acc*100:.2f}% (轮次 {self.best_epoch+1})")
        
        # 保存训练历史
        history = {
            'best_val_acc': self.best_val_acc,
            'best_epoch': self.best_epoch,
            'total_time': total_time,
            'metrics_history': self.metrics_history
        }
        
        with open(os.path.join(self.save_dir, 'training_history.json'), 'w', encoding='utf-8') as f:
            json.dump(history, f, indent=4, ensure_ascii=False)
        
        # 绘制最终的训练曲线
        self.plot_metrics()
        print(f"\n训练记录已保存到: {self.save_dir}")

    def train(self, model, train_loader, val_loader, criterion, optimizer, scheduler, num_epochs, device, classes):
        """训练模型
        
        参数:
            model: 要训练的模型
            train_loader: 训练数据加载器
            val_loader: 验证数据加载器
            criterion: 损失函数
            optimizer: 优化器
            scheduler: 学习率调度器
            num_epochs: 训练轮数
            device: 运行设备
            classes: 类别列表
        """
        self.epochs = num_epochs
        best_val_acc = 0
        start_time = time.time()
        
        try:
            for epoch in range(1, num_epochs + 1):
                # 开始新的训练轮次
                self.start_epoch(epoch, optimizer.param_groups[0]['lr'])
                
                # 训练阶段
                model.train()
                train_loss = 0
                train_correct = 0
                train_total = 0
                
                train_pbar = tqdm(train_loader, desc="训练进度")
                for inputs, labels in train_pbar:
                    inputs, labels = inputs.to(device), labels.to(device)
                    
                    # 清除梯度
                    optimizer.zero_grad()
                    
                    # 前向传播
                    outputs = model(inputs)
                    loss = criterion(outputs, labels)
                    
                    # 反向传播和优化
                    loss.backward()
                    optimizer.step()
                    
                    # 统计
                    train_loss += loss.item()
                    _, predicted = outputs.max(1)
                    train_total += labels.size(0)
                    train_correct += predicted.eq(labels).sum().item()
                    
                    # 更新进度条
                    train_acc = 100. * train_correct / train_total
                    train_pbar.set_postfix({
                        '损失': f'{train_loss/train_total:.4f}',
                        '准确率': f'{train_acc:.2f}%'
                    })
                    
                    # 显示GPU内存使用情况
                    if torch.cuda.is_available():
                        allocated = torch.cuda.memory_allocated() / 1024**3
                        reserved = torch.cuda.memory_reserved() / 1024**3
                        self.update_gpu_stats(allocated, reserved)
                
                train_pbar.close()
                
                # 计算训练指标
                train_loss = train_loss / len(train_loader)
                train_acc = 100. * train_correct / train_total
                
                # 验证阶段
                model.eval()
                val_loss = 0
                val_correct = 0
                val_total = 0
                
                with torch.no_grad():
                    val_pbar = tqdm(val_loader, desc="验证进度")
                    for inputs, labels in val_pbar:
                        inputs, labels = inputs.to(device), labels.to(device)
                        outputs = model(inputs)
                        loss = criterion(outputs, labels)
                        
                        val_loss += loss.item()
                        _, predicted = outputs.max(1)
                        val_total += labels.size(0)
                        val_correct += predicted.eq(labels).sum().item()
                        
                        # 更新进度条
                        val_acc = 100. * val_correct / val_total
                        val_pbar.set_postfix({
                            '损失': f'{val_loss/val_total:.4f}',
                            '准确率': f'{val_acc:.2f}%'
                        })
                    
                    val_pbar.close()
                
                # 计算验证指标
                val_loss = val_loss / len(val_loader)
                val_acc = 100. * val_correct / val_total
                
                # 更新学习率
                if isinstance(scheduler, optim.lr_scheduler.ReduceLROnPlateau):
                    scheduler.step(val_acc)
                else:
                    scheduler.step()
                
                # 保存最佳模型
                is_best = val_acc > best_val_acc
                if is_best:
                    best_val_acc = val_acc
                    self.best_val_acc = val_acc  # 更新类变量
                    self.best_epoch = epoch
                    self.save_checkpoint(model, optimizer, scheduler, is_best=True, val_acc=val_acc)  # 传入当前验证准确率
                
                # 显示本轮结果
                metrics = {
                    'train_loss': train_loss,
                    'train_acc': train_acc,
                    'val_loss': val_loss,
                    'val_acc': val_acc,
                    'best_acc': best_val_acc
                }
                self.end_epoch(metrics)
                
                # 每隔一定轮次绘制训练曲线
                if epoch % self.plot_interval == 0:
                    self.plot_training_curves()
                
                # 清理GPU内存
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                    gc.collect()
            
            # 训练结束，保存最终模型
            self.save_checkpoint(model, optimizer, scheduler, is_best=False, val_acc=val_acc)
            self.plot_training_curves()
            
            # 显示训练总结
            total_time = time.time() - start_time
            print(f"\n{Fore.GREEN}训练完成！{Style.RESET_ALL}")
            print(f"总用时: {total_time/3600:.2f}小时")
            print(f"最佳验证准确率: {best_val_acc:.2f}%（第{self.best_epoch}轮）")
            
        except KeyboardInterrupt:
            print(f"\n{Fore.YELLOW}训练被用户中断{Style.RESET_ALL}")
            self.save_checkpoint(model, optimizer, scheduler, is_best=False, val_acc=val_acc, interrupted=True)
            
        except Exception as e:
            print(f"\n{Fore.RED}训练过程出错: {str(e)}{Style.RESET_ALL}")
            raise
    
    def save_checkpoint(self, model, optimizer, scheduler, is_best=False, val_acc=None, interrupted=False):
        """保存检查点
        
        参数:
            model: 模型
            optimizer: 优化器
            scheduler: 学习率调度器
            is_best: 是否是最佳模型
            val_acc: 当前验证准确率
            interrupted: 是否是因中断而保存
        """
        checkpoint = {
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'scheduler_state_dict': scheduler.state_dict() if scheduler else None,
            'best_acc': self.best_val_acc,
            'epoch': self.best_epoch
        }
        
        # 保存最新检查点
        latest_path = os.path.join(self.save_dir, 'latest_checkpoint.pth')
        torch.save(checkpoint, latest_path)
        
        # 更新latest目录
        latest_checkpoint = os.path.join(self.latest_dir, 'latest_checkpoint.pth')
        shutil.copy2(latest_path, latest_checkpoint)
        
        # 如果是最佳模型，保存到多个位置
        if is_best and val_acc is not None:
            # 保存到实验目录
            best_path = os.path.join(self.save_dir, 'best_model.pth')
            torch.save(checkpoint, best_path)
            
            # 保存到model目录
            model_path = os.path.join(self.model_dir, 'best_model.pth')
            torch.save(checkpoint, model_path)
            
            # 更新latest目录的最佳模型
            latest_best = os.path.join(self.latest_dir, 'best_model.pth')
            shutil.copy2(best_path, latest_best)
            
            print(f"\n{Fore.GREEN}✨ 保存最佳模型，准确率: {val_acc:.2f}%")
            print(f"模型已保存到:")
            print(f"- 实验目录: {best_path}")
            print(f"- 模型目录: {model_path}")
            print(f"- 最新目录: {latest_best}{Style.RESET_ALL}")
        
        # 如果是中断保存，创建一个新的保存点
        if interrupted:
            interrupt_path = os.path.join(self.save_dir, 'interrupted_checkpoint.pth')
            torch.save(checkpoint, interrupt_path)
            
            # 更新latest目录的中断检查点
            latest_interrupt = os.path.join(self.latest_dir, 'interrupted_checkpoint.pth')
            shutil.copy2(interrupt_path, latest_interrupt)
            
            acc_info = f"，当前准确率: {val_acc:.2f}%" if val_acc is not None else ""
            print(f"\n{Fore.YELLOW}已保存中断检查点{acc_info}")
            print(f"- 实验目录: {interrupt_path}")
            print(f"- 最新目录: {latest_interrupt}{Style.RESET_ALL}")
    
    def plot_training_curves(self):
        """绘制训练曲线"""
        if not hasattr(self, 'history'):
            return
            
        plt.figure(figsize=(12, 8))
        
        # 绘制损失曲线
        plt.subplot(2, 1, 1)
        plt.plot(self.history['train_loss'], label='训练损失')
        plt.plot(self.history['val_loss'], label='验证损失')
        plt.title('训练过程中的损失变化')
        plt.xlabel('轮次')
        plt.ylabel('损失')
        plt.legend()
        plt.grid(True)
        
        # 绘制准确率曲线
        plt.subplot(2, 1, 2)
        plt.plot(self.history['train_acc'], label='训练准确率')
        plt.plot(self.history['val_acc'], label='验证准确率')
        plt.title('训练过程中的准确率变化')
        plt.xlabel('轮次')
        plt.ylabel('准确率 (%)')
        plt.legend()
        plt.grid(True)
        
        plt.tight_layout()
        plt.savefig(os.path.join(self.save_dir, 'training_curves.png'))
        plt.close()
    
    def log(self, message):
        """记录普通日志"""
        print(f"{message}")
    
    def log_warning(self, message):
        """记录警告信息"""
        print(f"{Fore.YELLOW}{message}{Style.RESET_ALL}")
    
    def log_error(self, message):
        """记录错误信息"""
        print(f"{Fore.RED}{message}{Style.RESET_ALL}")
